import random
import time

class ChessAI:
    def __init__(self, depth=2):
        self.depth = depth
        self.piece_values = {
            'pawn': 100,
            'knight': 320,
            'bishop': 330,
            'rook': 500,
            'queen': 900,
            'king': 20000
        }
        
        # 位置评估表，用于鼓励棋子控制中心和发展
        self.pawn_table = [
            [ 0,  0,  0,  0,  0,  0,  0,  0],
            [50, 50, 50, 50, 50, 50, 50, 50],
            [10, 10, 20, 30, 30, 20, 10, 10],
            [ 5,  5, 10, 25, 25, 10,  5,  5],
            [ 0,  0,  0, 20, 20,  0,  0,  0],
            [ 5, -5,-10,  0,  0,-10, -5,  5],
            [ 5, 10, 10,-20,-20, 10, 10,  5],
            [ 0,  0,  0,  0,  0,  0,  0,  0]
        ]
        
        self.knight_table = [
            [-50,-40,-30,-30,-30,-30,-40,-50],
            [-40,-20,  0,  0,  0,  0,-20,-40],
            [-30,  0, 10, 15, 15, 10,  0,-30],
            [-30,  5, 15, 20, 20, 15,  5,-30],
            [-30,  0, 15, 20, 20, 15,  0,-30],
            [-30,  5, 10, 15, 15, 10,  5,-30],
            [-40,-20,  0,  5,  5,  0,-20,-40],
            [-50,-40,-30,-30,-30,-30,-40,-50]
        ]
        
        self.bishop_table = [
            [-20,-10,-10,-10,-10,-10,-10,-20],
            [-10,  0,  0,  0,  0,  0,  0,-10],
            [-10,  0, 10, 10, 10, 10,  0,-10],
            [-10,  5,  5, 10, 10,  5,  5,-10],
            [-10,  0,  5, 10, 10,  5,  0,-10],
            [-10,  5,  5,  5,  5,  5,  5,-10],
            [-10,  0,  5,  0,  0,  5,  0,-10],
            [-20,-10,-10,-10,-10,-10,-10,-20]
        ]
        
        self.rook_table = [
            [ 0,  0,  0,  0,  0,  0,  0,  0],
            [ 5, 10, 10, 10, 10, 10, 10,  5],
            [-5,  0,  0,  0,  0,  0,  0, -5],
            [-5,  0,  0,  0,  0,  0,  0, -5],
            [-5,  0,  0,  0,  0,  0,  0, -5],
            [-5,  0,  0,  0,  0,  0,  0, -5],
            [-5,  0,  0,  0,  0,  0,  0, -5],
            [ 0,  0,  0,  5,  5,  0,  0,  0]
        ]
        
        self.queen_table = [
            [-20,-10,-10, -5, -5,-10,-10,-20],
            [-10,  0,  0,  0,  0,  0,  0,-10],
            [-10,  0,  5,  5,  5,  5,  0,-10],
            [ -5,  0,  5,  5,  5,  5,  0, -5],
            [  0,  0,  5,  5,  5,  5,  0, -5],
            [-10,  5,  5,  5,  5,  5,  0,-10],
            [-10,  0,  5,  0,  0,  0,  0,-10],
            [-20,-10,-10, -5, -5,-10,-10,-20]
        ]
        
        self.king_middle_table = [
            [-30,-40,-40,-50,-50,-40,-40,-30],
            [-30,-40,-40,-50,-50,-40,-40,-30],
            [-30,-40,-40,-50,-50,-40,-40,-30],
            [-30,-40,-40,-50,-50,-40,-40,-30],
            [-20,-30,-30,-40,-40,-30,-30,-20],
            [-10,-20,-20,-20,-20,-20,-20,-10],
            [ 20, 20,  0,  0,  0,  0, 20, 20],
            [ 20, 30, 10,  0,  0, 10, 30, 20]
        ]
        
        self.king_end_table = [
            [-50,-40,-30,-20,-20,-30,-40,-50],
            [-30,-20,-10,  0,  0,-10,-20,-30],
            [-30,-10, 20, 30, 30, 20,-10,-30],
            [-30,-10, 30, 40, 40, 30,-10,-30],
            [-30,-10, 30, 40, 40, 30,-10,-30],
            [-30,-10, 20, 30, 30, 20,-10,-30],
            [-30,-30,  0,  0,  0,  0,-30,-30],
            [-50,-30,-30,-30,-30,-30,-30,-50]
        ]
    
    def get_best_move(self, board, color):
        """获取最佳移动"""
        start_time = time.time()
        best_move = None
        best_score = float('-inf')
        
        # 收集所有可能的移动
        all_moves = []
        for row in range(8):
            for col in range(8):
                piece = board.board[row][col]
                if piece and piece.color == color:
                    valid_moves = board.get_valid_moves(row, col)
                    for move in valid_moves:
                        all_moves.append((row, col, move[0], move[1]))
        
        # 如果没有可用移动，返回None
        if not all_moves:
            return None
            
        # 随机打乱移动顺序，增加变化性
        random.shuffle(all_moves)
        
        # 使用极小极大算法评估每个移动
        for move in all_moves:
            src_row, src_col, dst_row, dst_col = move
            
            # 模拟移动
            piece = board.board[src_row][src_col]
            captured_piece = board.board[dst_row][dst_col]
            
            board.board[dst_row][dst_col] = piece
            board.board[src_row][src_col] = None
            
            # 如果移动的是国王，更新其位置
            original_king_pos = None
            if piece.type == 'king':
                original_king_pos = board.king_positions[color]
                board.king_positions[color] = (dst_row, dst_col)
            
            # 评估移动后的局面
            score = -self._minimax(board, self.depth - 1, float('-inf'), float('inf'), 
                                  'white' if color == 'black' else 'black')
            
            # 恢复棋盘状态
            board.board[src_row][src_col] = piece
            board.board[dst_row][dst_col] = captured_piece
            
            # 恢复国王位置
            if original_king_pos:
                board.king_positions[color] = original_king_pos
            
            # 更新最佳移动
            if score > best_score:
                best_score = score
                best_move = move
            
            # 超过3秒，提前结束搜索
            if time.time() - start_time > 3:
                break
        
        return best_move
    
    def _minimax(self, board, depth, alpha, beta, color):
        """极小极大算法带Alpha-Beta剪枝"""
        if depth == 0:
            return self._evaluate_board(board)
        
        # 检查将军或逼和
        opponent_color = 'white' if color == 'black' else 'black'
        if board.is_checkmate(color):
            return 10000 if color == 'black' else -10000  # 对手被将死
        elif board.is_stalemate(color):
            return 0  # 和棋
        
        best_score = float('-inf') if color == 'white' else float('inf')
        
        # 收集当前玩家所有可能的移动
        all_moves = []
        for row in range(8):
            for col in range(8):
                piece = board.board[row][col]
                if piece and piece.color == color:
                    valid_moves = board.get_valid_moves(row, col)
                    for move in valid_moves:
                        all_moves.append((row, col, move[0], move[1]))
        
        # 随机打乱移动顺序，增加变化
        random.shuffle(all_moves)
        
        # 评估每个移动
        for move in all_moves:
            src_row, src_col, dst_row, dst_col = move
            
            # 模拟移动
            piece = board.board[src_row][src_col]
            captured_piece = board.board[dst_row][dst_col]
            
            board.board[dst_row][dst_col] = piece
            board.board[src_row][src_col] = None
            
            # 如果移动的是国王，更新其位置
            original_king_pos = None
            if piece.type == 'king':
                original_king_pos = board.king_positions[color]
                board.king_positions[color] = (dst_row, dst_col)
            
            # 递归评估
            score = self._minimax(board, depth - 1, alpha, beta, 
                                 'white' if color == 'black' else 'black')
            
            # 恢复棋盘状态
            board.board[src_row][src_col] = piece
            board.board[dst_row][dst_col] = captured_piece
            
            # 恢复国王位置
            if original_king_pos:
                board.king_positions[color] = original_king_pos
            
            # 更新最佳分数和Alpha-Beta值
            if color == 'white':
                best_score = max(best_score, score)
                alpha = max(alpha, best_score)
            else:
                best_score = min(best_score, score)
                beta = min(beta, best_score)
            
            # Alpha-Beta剪枝
            if beta <= alpha:
                break
        
        return best_score
    
    def _evaluate_board(self, board):
        """评估棋盘局面"""
        total_score = 0
        
        # 统计所有棋子，同时判断是否进入残局
        piece_count = {'white': 0, 'black': 0}
        
        for row in range(8):
            for col in range(8):
                piece = board.board[row][col]
                if piece:
                    # 计算棋子基础分值
                    piece_value = self.piece_values[piece.type]
                    
                    # 根据棋子位置增加额外评分
                    position_value = 0
                    if piece.type == 'pawn':
                        position_value = self.pawn_table[row][col]
                    elif piece.type == 'knight':
                        position_value = self.knight_table[row][col]
                    elif piece.type == 'bishop':
                        position_value = self.bishop_table[row][col]
                    elif piece.type == 'rook':
                        position_value = self.rook_table[row][col]
                    elif piece.type == 'queen':
                        position_value = self.queen_table[row][col]
                    elif piece.type == 'king':
                        # 根据是否残局选择不同的国王评估表
                        is_endgame = self._is_endgame(board)
                        if is_endgame:
                            position_value = self.king_end_table[row][col]
                        else:
                            position_value = self.king_middle_table[row][col]
                    
                    # 计算总分
                    score = piece_value + position_value
                    
                    # 根据棋子颜色决定加减分
                    if piece.color == 'white':
                        total_score += score
                        piece_count['white'] += 1
                    else:
                        total_score -= score
                        piece_count['black'] += 1
        
        return total_score
    
    def _is_endgame(self, board):
        """检查是否进入残局阶段"""
        # 当双方都没有皇后或者双方棋子少于10个时算作残局
        white_queen = black_queen = False
        piece_count = 0
        
        for row in range(8):
            for col in range(8):
                piece = board.board[row][col]
                if piece:
                    piece_count += 1
                    if piece.type == 'queen':
                        if piece.color == 'white':
                            white_queen = True
                        else:
                            black_queen = True
        
        return (not white_queen and not black_queen) or piece_count <= 10 